"""UNet++训练脚本"""
import os, sys
sys.path.append(os.path.join(os.path.dirname(__file__), '../..'))

import argparse
from baselines.unetpp.model import UNetPlusPlus
from baselines.unetpp import config as config_module
from train_template import train_model
import json, numpy as np

def main():
    parser = argparse.ArgumentParser(description='UNet++训练脚本')
    parser.add_argument('--fold', type=int, default=None)
    parser.add_argument('--resume', action='store_true')
    parser.add_argument('--epochs', type=int, default=None)
    parser.add_argument('--batch_size', type=int, default=None)
    parser.add_argument('--lr', type=float, default=None)
    args = parser.parse_args()

    if args.fold is not None:
        best_metrics = train_model('unetpp', UNetPlusPlus, config_module, args.fold, args)
        print(f'Fold {args.fold} 最佳指标: {best_metrics}')
    else:
        all_results = {}
        for fold_idx in range(config_module.TRAIN_CONFIG['k_folds']):
            print(f'\n{"=" * 80}\n开始训练 Fold {fold_idx}\n{"=" * 80}\n')
            best_metrics = train_model('unetpp', UNetPlusPlus, config_module, fold_idx, args)
            all_results[f'fold_{fold_idx}'] = best_metrics

        metrics_names = list(all_results['fold_0'].keys())
        summary = {}
        for metric_name in metrics_names:
            values = [all_results[f'fold_{i}'][metric_name] for i in range(config_module.TRAIN_CONFIG['k_folds'])]
            summary[metric_name] = {'mean': np.mean(values), 'std': np.std(values), 'values': values}

        for metric_name, stats in summary.items():
            print(f'{metric_name}: Mean: {stats["mean"]:.4f}, Std: {stats["std"]:.4f}')

        summary_path = os.path.join('logs', 'unetpp', 'summary_results.json')
        os.makedirs(os.path.dirname(summary_path), exist_ok=True)
        with open(summary_path, 'w', encoding='utf-8') as f:
            json.dump({'all_results': all_results, 'summary': summary}, f, indent=4, ensure_ascii=False)

if __name__ == '__main__':
    main()
